import cv2
import numpy as np
import onnxruntime as ort
import torch
from icecream import ic
# from PIL import Image


def img_preprocess(img, new_shape=(112, 112)):
    # print(img.shape)
    img = cv2.resize(img, new_shape)
    # ic(img.shape)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # ic(img.shape)
    img = np.transpose(img, (2, 0, 1))
    # ic(img.shape)
    img = torch.from_numpy(img).unsqueeze(0).float()
    # ic(img.shape)
    img = img.numpy()
    # ic(img.shape)
    img = (img - 127.5) / 127.5
    # ic(img.shape)
    # img.div_(255).sub_(0.5).div_(0.5)

    return img


class Arcface:
    def __init__(self, onnxpath):
        cuda = torch.cuda.is_available()
        providers = (
            ["CUDAExecutionProvider", "CPUExecutionProvider"]
            if cuda
            else ["CPUExecutionProvider"]
        )
        self.onnx_session = ort.InferenceSession(onnxpath, providers=providers)

        self.input_name = self.get_input_name()
        self.output_name = self.get_output_name()
        self.new_shape = (112, 112)

    def get_input_name(self):
        input_name = []
        for node in self.onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_output_name(self):
        output_name = []
        for node in self.onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def predict_img_nms(self, img_path):
        """
        list
        (0, 0.954, ((173, 155), (432, 253)))
        cls,thresh,x1,y1,x2,y2
        """
        # print('input',self.input_name)
        # print('output',self.output_name)
        img = cv2.imread(img_path)
        im = img_preprocess(img, self.new_shape)
        inp = {self.input_name[0]: im}
        outputs = self.onnx_session.run(self.output_name, inp)
        # pred = torch.from_numpy(outputs)
        # print(outputs)
        return outputs[0][0]

    def predict_image(self, image):
        im = img_preprocess(image, self.new_shape)
        inp = {self.input_name[0]: im}
        outputs = self.onnx_session.run(self.output_name, inp)
        return outputs[0][0]

    def cal_sim(self, feat1, feat2):
        feat1 = feat1 / np.linalg.norm(feat1)
        feat2 = feat2 / np.linalg.norm(feat2)
        cur_score = np.dot(feat1, feat2)  # 余弦相似度（即得分）
        return cur_score


if __name__ == "__main__":
    onnx_path = "./models/new/face_feature_230320.onnx"

    img1_path = "data/Aaron_Guiel_0001.jpg"
    img2_path = "data/Aaron_Guiel_0002.jpg"
    
    # img1_path = "data/Abdullah_0002.jpg"
    # img2_path = "data/Abdullah_0003.jpg"

    model = Arcface(onnx_path)
    fea1 = model.predict_img_nms(img1_path)
    ic(fea1.shape)
    fea2 = model.predict_img_nms(img2_path)
    cur_score = model.cal_sim(fea1, fea2)
    print(cur_score)
